import unittest

from langchain_core.messages import SystemMessage, HumanMessage
from langchain_core.outputs import LLMResult

from llm_repeator_redis.llm_services import LLMResourcesManager


class MyTestCase(unittest.TestCase):
    def test_generate(self):
        # self.assertEqual(True, False)  # add assertion here
        llm_manager = LLMResourcesManager('../config/llm_resources.json')

        messages = [
            SystemMessage("请你使用中文回答我的问候"),
            HumanMessage("How are you?")
        ]

        for llm_id in llm_manager.list_llm_def():
            print(llm_id)

            content: LLMResult = llm_manager.generate(llm_key=llm_id, base_messages=messages)
            print(content)

            self.assertIsNotNone(content)

    def test_stream(self):
        # self.assertEqual(True, False)  # add assertion here
        llm_manager = LLMResourcesManager('../config/llm_resources.json')

        messages = [
            SystemMessage("请你使用中文回答我的问候"),
            HumanMessage("How are you?")
        ]

        print("\n\nchat_1")
        for chunk in llm_manager.stream(llm_key="huawei_deepseek_v3_32k", base_messages=messages):
            print(chunk.content, end="", flush=True)
            self.assertIsNotNone(chunk)

        print("\n\nchat_2")
        for chunk in llm_manager.stream(llm_key="huawei_deepseek_r1_32k", base_messages=messages):
            print(chunk.content, end="", flush=True)
            self.assertIsNotNone(chunk)

        print("\n\nchat_3")
        for chunk in llm_manager.stream(llm_key="home_qwen3:32b", base_messages=messages):
            print(chunk, end="", flush=True)
            self.assertIsNotNone(chunk)




if __name__ == '__main__':
    unittest.main()
